#! /usr/bin/env python3
import cv2
import numpy as np
from PIL import Image
import os
import argparse
from tqdm import tqdm

def active(X, W):
    M = W * (1 - X.astype('float32') / 255)
    C = 1 - M * (M > 0.25)
    A = np.tan(np.pi / 3 * C - np.pi / 12)
    return (np.maximum(A, 0) * 255).astype('uint8')

def textCleaner(input_image, output_path):
    if isinstance(input_image, Image.Image):
        img = np.array(input_image.convert('L'), dtype='uint8')
    else:
        img = cv2.imread(input_image, 0)

    h, w = img.shape[:2]
    k = 600 / max(h, w)
    mini = cv2.resize(img, (round(w / 4), round(h / 4)))

    kernel = np.ones((3, 3))

    edges = cv2.Canny(mini, 100, 200)

    morphology = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel, iterations=1)
    median = cv2.medianBlur(morphology, 3)

    morphology = cv2.morphologyEx(median, cv2.MORPH_OPEN, kernel, iterations=1)
    morphology = cv2.morphologyEx(morphology, cv2.MORPH_DILATE, kernel, iterations=2)

    kernel = np.ones((5, 5)) / 9
    filted = cv2.filter2D(morphology, -1, kernel)

    mask = filted.astype('float32') / filted.max()
    mask = cv2.resize(mask, (w, h))
    clean = active(img, mask)
    cv2.imwrite(output_path, clean)

def textCleaners(input_dir, output_dir):
    basenames = os.listdir(input_dir)
    for basename in tqdm(basenames, ncols=80):
        input_path = os.path.join(input_dir, basename)
        output_path = os.path.join(output_dir, basename)
        textCleaner(input_path, output_path)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='背景除去，图片抗噪，信息提升工具')
    parser.add_argument('input_dir', metavar='原始图片目录', type=str,
                        help='原始图片所在的目录路径')
    parser.add_argument('output_dir', metavar='生成图片目录', type=str,
                        help='生成图片所在的目录路径')

    args = parser.parse_args()
    textCleaners(args.input_dir, args.output_dir)
